# -*- coding: utf-8 -*-
#
#    imap.py
#       
#    Copyright 2012 Jose A. Jimenez <jajimc@gmail.com>
#       
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Unimail
@version: 0.14
@license: GPL 3.0 (see "license" file)
@author: Jose A. Jimenez
@contact: jajimc@gmail.com"""

import os
import email
import imaplib
import threading
import signal
import excep
        
def get_response(imapObj):
    """Modified function from 'imaplib.IMAP4' class to allow cancelation of it.
    @param imapObj: IMAP4 object
    @type imapObj: L{imaplib.IMAP4}"""
    resp = imapObj._get_line()

    if imapObj._match(imapObj.tagre, resp):
        tag = imapObj.mo.group('tag')
        if not tag in imapObj.tagged_commands:
            raise imapObj.abort('unexpected tagged response: %s' % resp)
        typ = imapObj.mo.group('type')
        dat = imapObj.mo.group('data')
        imapObj.tagged_commands[tag] = (typ, [dat])
    else:
        dat2 = None

        if not imapObj._match(imaplib.Untagged_response, resp):
            if imapObj._match(imaplib.Untagged_status, resp):
                dat2 = imapObj.mo.group('data2')

        if imapObj.mo is None:
            if imapObj._match(imaplib.Continuation, resp):
                imapObj.continuation_response = imapObj.mo.group('data')
                return None
            raise imapObj.abort("unexpected response: '%s'" % resp)

        typ = imapObj.mo.group('type')
        dat = imapObj.mo.group('data')
        if dat is None: dat = ''
        if dat2: dat = dat + ' ' + dat2

        while imapObj._match(imaplib.Literal, dat) and imapObj.doProcess:
            size = int(imapObj.mo.group('size'))
            if __debug__:
                if imapObj.debug >= 4:
                    imapObj._mesg('read literal size %s' % size)
            data = imapObj.read(size)
            imapObj._append_untagged(typ, (dat, data))
            dat = imapObj._get_line()

        imapObj._append_untagged(typ, dat)

    if typ in ('OK', 'NO', 'BAD') and imapObj._match(imaplib.Response_code, dat):
        imapObj._append_untagged(imapObj.mo.group('type'), imapObj.mo.group('data'))

    if __debug__:
        if imapObj.debug >= 1 and typ in ('NO', 'BAD', 'BYE'):
            imapObj._mesg('%s response: %s' % (typ, dat))

    return resp
        
class Imap4(imaplib.IMAP4):
    """Like L{imaplib.IMAP4} but allows to cancel the process and set the maximum download rate."""

    def __init__(self, host, port, maxDownloadRate=1024):
        """@param host: Imap host address
        @type host: C{str}
        @param port: Imap port
        @type port: C{int}
        @param maxDownloadRate: Maximum download rate in Kbs. Must be at least 1024.
        @type maxDownloadRate: C{int}"""
        if type(maxDownloadRate) != int or maxDownloadRate < 1024: raise Exception("'maxDownloadRate' must be an integer greater or equal than 1024.")
        
        self.__maxDownloadRate = maxDownloadRate
        imaplib.IMAP4.__init__(self, host, port)

    def read(self, size):
        """Redefinition of L{imaplib.IMAP4.read} method."""
        return imaplib.IMAP4.read(self.__maxDownloadRate)
        
    def fetch(self, message_set, message_parts):
        """Redefinition of L{imaplib.IMAP4.fetch} method."""
        self.doProcess = True
        imaplib.IMAP4.fetch(self, message_set, message_parts)
        
    def _get_response(self):
        """Redefinition of L{imaplib.IMAP4._get_response} method."""
        return get_response(self)
    
    def cancelProcess(self):
        """Cancels current download process."""
        self.doProcess = False
        
class Imap4Ssl(imaplib.IMAP4_SSL):
    """Like L{imaplib.IMAP4_SSL} but allows to cancel the process and set the maximum download rate."""
    
    def __init__(self, host, port, maxDownloadRate=1024):
        """@param host: Imap host address
        @type host: C{str}
        @param port: Imap port
        @type port: C{int}
        @param maxDownloadRate: Maximum download rate in Kbs. Must be at least 1024.
        @type maxDownloadRate: C{int}"""
        if type(maxDownloadRate) != int or maxDownloadRate < 1024: raise Exception("'maxDownloadRate' must be an integer greater or equal than 1024.")
        
        self.__maxDownloadRate = maxDownloadRate
        imaplib.IMAP4_SSL.__init__(self, host, port)

    def read(self, size):
        """Redefinition of L{imaplib.IMAP4_SSL.read} method."""
        return imaplib.IMAP4_SSL.read(self.__maxDownloadRate)
        
    def fetch(self, message_set, message_parts):
        """Redefinition of L{imaplib.IMAP4_SSL.fetch} method."""
        self.doProcess = True
        imaplib.IMAP4.fetch(self, message_set, message_parts)
        
    def _get_response(self):
        """Redefinition of L{imaplib.IMAP4_SSL._get_response} method."""
        return get_response(self)
    
    def cancelProcess(self):
        """Cancels current download process."""
        self.doProcess = False
        
class EmailAttachment:
    """Represents an e-mail message attached file."""

    def __init__(self, content, name):
        """@param content: File content
        @type content: C{str}
        @param fileName: File name
        @type fileName: C{str}"""
        if type(content) != str or content == "": raise Exception("'content' must be a non empty string.")
        if type(name) != str or name == "": raise Exception("'name' must be a non empty string.")
        
        self.__content = content
        self.__name = fix_text(name)
    
    def save(self, dirPath, fileName=None):
        """Saves the attachment to a file.
        @param dirPath: Parent directory where to save the attachment
        @type dirPath: C{str}
        @param fileName: Output file name. If omitted, its value is the one given by the class constructor.
        @type fileName: C{str}"""
        if type(dirPath) != str or dirPath == "": raise Exception("'dirPath' must be a non empty string.")
        if type(fileName) != None and (type(fileName) != str or fileName == ""): raise Exception("'fileName' must be a non empty string or None.")
        if not os.path.isdir(dirPath): raise Exception("'dirPath' must be an existing directory.")
        
        try:
            if (fileName == None):
                path = os.path.join(dirPath, self.__fileName)
            else:
                path = os.path.join(dirPath, fileName)
                
            outputFile = open(path, "wb")
            outputFile.write(self.__content)
        except Exception as exc:
            raise exc
        finally:
            outputFile.close()
        
class EmailAttachmentSet:
    """Set of L{EmailAttachment} of an e-mail message."""

    def __init__(self, messageData):
        """@param message: E-mail message data
        @type message: C{list}"""
        if type(messageData) != list or len(messageData) < 1: raise Exception("'messageData' must be a list of C{str} with at least 1 element.")
        
        self.__attachmentsIndex = 0 
        self.__attachments = []
    
        for part in messageData.walk():
            fileName = part.get_filename()
            if fileName != None:
                fileContent = part.get_payload(decode=1)
                self.__attachments.append(EmailAttachment(fileContent, fileName))
    
    def count(self):
        """Returns the number of attachments of the set.
        @return: Attachment count
        @rtype: C{int}"""
        return len(self.__attachments)
        
    def __iter__(self):
        """Mandatory function in order to make the class iterable.
        @returns: Current object
        @rtype: L{EmailAttachmentSet}"""
        return self

    def next(self):
        """Mandatory function in order to make the class iterable. Returns next attachment.
        @returns: Next attachment
        @rtype: L{EmailAttachment}"""

        try: 
            attachment = self.__attachments[self.__attachmentsIndex]
        except:
            raise StopIteration
            
        self.__attachmentsIndex += 1
        return attachment
        
class EmailMessagePreview:
    """Represents an e-mail message."""

    def __init__(self, index, data):
        """@param index: Message inbox index. Must be greater or equal than 1.
        @type index: C{int}
        @param data: Message contents
        @type data: C{str}"""
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")
        if type(data) != data or data == "": raise Exception("'data' must be a non empty string.")
        
        self.__index = index
        content = email.message_from_string(data[0][1])

        self.__subject = content["subject"]
        self.__fromAddress = content["from"]
        
    def get_index(self):
        """return: Message index
        @rtype: C{int}"""
        return self.__index
        
    def get_subject(self):
        """return: Message subject
        @rtype: C{str}"""
        return self.__subject
        
    def get_from_address(self):
        """return: Message subject
        @rtype: C{str}"""
        return self.__fromAddress
    
    def __cmp__(self, other):
        """Compares two objets of this class by their subjects.
        @return: -1 if current object is first order, 1 if the other is first order or 0 if both are equal
        @rtype: C{int}"""
        if self.__subject == other.__subject:
            result = 0
        elif self.__subject < other.__subject:
            result = -1
        else:
            result = 1
        return result
        
class EmailMessage():
    """Represents a full e-mail message (with attachments)."""

    def __init__(self, index, data):
        """@param index: Message inbox index. Must be equal or greater than 1.
        @type index: C{int}
        @param data: Message contents
        @type data: C{str}"""
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")
        if type(data) != data or data == "": raise Exception("'data' must be a non empty string.")
        
        self.__index = index
        content = email.message_from_string(data[0][1])
        
        self.__subject = content["subject"]
        self.__fromAddress = content["from"]
        self.__attachments = EmailAttachmentSet(content)
    
    def get_index(self):
        """return: Message index
        @rtype: C{int}"""
        return self.__index
        
    def get_subject(self):
        """return: Message subject
        @rtype: C{str}"""
        return self.__subject
        
    def get_from_address(self):
        """return: Message subject
        @rtype: C{str}"""
        return self.__fromAddress
       
    def get_attachments(self):
        """return: Message attachments
        @rtype: L{EmailAttachmentSet}"""
        return self.__attachments
    
    def __cmp__(self, other):
        """Compares two objets of this class by their subjects.
        @return: -1 if current object is first order, 1 if the other is first order or 0 if both are equal
        @rtype: C{int}"""
        if self.__subject == other.__subject:
            result = 0
        elif self.__subject < other.__subject:
            result = -1
        else:
            result = 1
        return result
        
class ImapClient():
    """IMAP4 client class."""
    __connection = None
    __messageCount = None
    
    def __init__(self, host, ssl, port, userId, userPassword):
        """@param host: IMAP host address
        @type host: C{str}
        @param ssl: Whether to use SSL or not
        @type ssl: C{bool}
        @param port: IMAP host port
        @type port: C{int}
        @param userId: User ID
        @type userId: C{str}
        @param userPassword: User password
        @type userPassword: C{str}"""
        self.set_host(host)
        self.set_ssl(ssl)
        self.set_port(port)
        self.set_user_id(userId)
        self.set_user_password(userPassword)
        
    def get_host(self):
        """@return: IMAP host address
        @rtype: C{str}"""
        return self.__host
    
    def get_ssl(self):
        """@return: Whether if SSL is enabled or not
        @rtype: C{bool}"""
        return self.__ssl
        
    def get_port(self):
        """@return: IMAP host port
        @rtype: C{int}"""
        return self.__port
        
    def get_user_id(self):
        """@return: User ID
        @rtype: C{str}"""
        return self.__userId
        
    def get_user_password(self):
        """@return: User password
        @rtype: C{str}"""
        return self.__userPassword
        
    def set_host(self, host):
        """Sets the IMAP host address.
        @param host: Host address
        @type host: C{str}"""
        if type(host) != str: raise Exception("'host' must be a string.")
        self.__host = host
        
    def set_ssl(self, ssl):
        """Enables/disables SSL.
        @param ssl: Whether if SSL is enabled or not
        @type ssl: C{bool}"""
        if type(ssl) != bool: raise Exception("'ssl' must be a boolean.")
        self.__ssl = ssl
        
    def set_port(self, port):
        """Sets the IMAP host port.
        @param port: Port. Must be an integer greater or equal than 0.
        @type port: C{int}"""
        if type(port) != int or port < 0: raise Exception("'port' must be an integer greater or equal than 0.")
        self.__port = port
        
    def set_user_id(self, id):
        """Sets the user ID.
        @param id: ID
        @type id: C{str}"""
        if type(id) != str or id == "": raise Exception("'id' must be a non empty string.")
        self.__userId = id
        
    def set_user_password(self, password):
        """Sets the user password.
        @param password: Password
        @type password: C{str}"""
        if type(password) != str or password == "": raise Exception("'password' must be a non empty string.")
        self.__userPassword = password
    
    def check_host(self, timeout=10):
        """Checks whether if host is available or not.
        @param timeout: Timeout in seconds to do the check. It must be greater or equal than 1. If omitted, its value is 10.
        @type timeout: C{int}"""
        if type(timeout) != int or timeout < 1: raise Exception("'timeout' must be an integer greater or equal than 1.")
        
        available = False
        signal.signal(signal.SIGALRM, excep.process_timeout)
        signal.alarm(timeout) 
        
        try:
            if self.__ssl:
                connection = imaplib.IMAP_SSL(self.__host, self.__port)
            else:
                connection = imaplib.IMAP(self.__host, self.__port)
            
            available = True
        except:
            pass
            
        signal.alarm(0)
        return available
    
    def check_credentials(self, timeout=10):
        """Checks whether if user and password are valid or not.
        @param timeout: Timeout in seconds to do the check. It must be greater or equal than 1. If omitted, its value is 10.
        @type timeout: C{int}"""
        if type(timeout) != int or timeout < 1: raise Exception("'timeout' must be an integer greater or equal than 1.")
        
        valid = False
        signal.signal(signal.SIGALRM, excep.process_timeout)
        signal.alarm(timeout)
        
        try:
            if self.__ssl:
                connection = imaplib.IMAP4_SSL(self.__host, self.__port)
            else:
                connection = imaplib.IMAP4(self.__host, self.__port)
        
            connection.login(self.__userId, self.__userPassword)
            valid = True
        except:
            pass
            
        signal.alarm(0)
        return valid
        
    def connect(self, timeout=10):
        """Connects to host.
        @param timeout: Timeout in seconds to do the check. It must be greater or equal than 1. If omitted, its value is 10.
        @type timeout: C{int}
        @raise excep.TimeoutException: If function takes more time than L{timeout}
        @raise excep.AuthenticationException: If user and password are not valid"""
        if type(timeout) != int or timeout < 1: raise Exception("'timeout' must be an integer greater or equal than 1.")
        
        valid = False
        signal.signal(signal.SIGALRM, excep.process_timeout)
        signal.alarm(timeout)
        
        try:
            if self.__ssl:
                self.__connection = Imap4Ssl(self.__host, self.__port, 1024)
            else:
                self.__connection = Imap4(self.__host, self.__port, 1024)
            self.__connection.login(self.__userId, self.__userPassword)
        except excep.TimeoutException as timeoutExc:
            raise timeoutExc
        except Exception as exc:
            raise excep.AuthenticationException()
            
        signal.alarm(0)
        
    def cancel_process(self):
        """Cancels current process."""
        if self.__connection != None:
            self.__connection.cancelProcess()
        
    def close(self):
        """Closes current open connection."""
        if self.__connection != None:
            self.__connection.logout()
            self.__connection = None
            
    def get_message_count(self, timeout=10):
        """Get inbox message count. L{connect} must be called first.
        @param timeout: Timeout in seconds to do the check. It must be greater or equal than 1. If omitted, its value is 10.
        @type timeout: C{int}
        @raise excep.TimeoutException: If function takes more time than L{timeout}
        @raise Exception: If some other error occurs"""
        if type(timeout) != int or timeout < 1: raise Exception("'timeout' must be an integer greater or equal than 1.")
        
        signal.signal(signal.SIGALRM, excep.process_timeout)
        signal.alarm(timeout)
        
        try:
            typ, data = self.__connection.select("INBOX")
            messageCount = int(data[0])
        except Exception as exc:
            raise exc
            
        signal.alarm(0)
        return messageCount
            
    def get_messages_by_subject(self, text, timeout=10):
        """Returns a list of indexes of each message whose subject contains L{text}.
        @param text: Substring to search in subjects
        @type text: C{str}
        @param timeout: Timeout in seconds to do the check. It must be greater or equal than 1. If omitted, its value is 10.
        @type timeout: C{int}
        @raise excep.TimeoutException: If function takes more time than L{timeout}
        @raise Exception: If some other error occurs"""
        if type(timeout) != int or timeout < 1: raise Exception("'timeout' must be an integer greater or equal than 1.")
        
        signal.signal(signal.SIGALRM, excep.process_timeout)
        signal.alarm(timeout)
        
        try:
            searchString = "(SUBJECT \"" + text + "\")"
            typ, [ids] = self.connection.search(None, searchString)
            idList = ids.split()
            for i in range(0, len(idList)):
                idList[i] = int(idList[i])
        except Exception as exc:
            raise exc
                
        signal.alarm(0)
        return idList
        
    def get_message_preview(self, index):
        """Gets a message given by its index.
        @param index: Message index. Must be greater or equal than 1.
        @type index: C{int}
        @return: Message whose index is L{index}
        @rtype: L{EmailMessagePreview}""" 
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")

        typ, data = self.__connection.fetch(str(index), "(BODY.PEEK[HEADER])")
        return EmailMessagePreview(index, data)
        
    def get_message(self, index):
        """Gets a full message (with attachments) given by its index.
        @param index: Message index. Must be greater or equal than 1.
        @type index: C{int}
        @return: Message whose index is L{index}
        @rtype: L{EmailMessage}""" 
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")

        typ, data = self.__connection.fetch(str(index), "(RFC822)")
        return EmailMessage(index, data)
        
    def delete_message(self, index):
        """Deletes a message by its index.
        @param index: Message index. Must be greater or equal than 1.
        @type index: C{int}"""
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")
        
        typ, response = self.__connection.store(str(index), '+FLAGS', r'(\Deleted)')
        typ, response = self.__connection.expunge()
        
class ImapDownloadThread(threading.Thread):
    #ToDo
    """Manages an e-mail message download."""
    __message = None
    
    def __init__(self, imapClient, messageIndex, errorHandler, progressHandler, completeHandler):
        """@param imapClient: L{ImapClient} object to use. Must have an open connection.
        @type imapClient: L{ImapClient}
        @param messageIndex: Message inbox index. Must be greater or equal than 1.
        @type messageIndex: C{int}
        @param errorHandler: Handler function to be called when an error occurs
        @type errorHandler: C{function}
        @param progressHandler: Handler function to be called when download progress is updated
        @type progressHandler: C{function}
        @param completeHandler: Handler function to be called when download is complete
        @type completeHandler: C{function}"""
        if not isinstance(imapClient, ImapClient) raise Exception("'imapClient' must be an 'ImapClient' object.")
        if type(index) != int or index < 1: raise Exception("'index' must be an integer greater or equal than 1.")
        
        self.__imapClient = imapClient
        self.__messageIndex = messageIndex
        self.__errorHandler = errorHandler
        self.__progressHandler = progressHandler
        self.__completeHandler = completeHandler
        
    def run(self):
        """Starts download"""
        #ToDo: Handle "error", "progress" and "complete" events
        self.__message = self.__imapClient.get_message(self.__messageIndex)
        
    def get_message(self):
        """Returns downloaded message
        @return: Downloaded message
        @rtype: L{EmailMessage}"""
        return self.__message
        
    def close(self):
        """Cancels download"""
        imapClient.cancel_process()
